﻿#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <ctime>
#include <iomanip>
#include <unordered_map>

#ifdef _WIN32
#define WIN32_LEAN_AND_MEAN
#include <winsock2.h>  //网络API
#include <windows.h>  //windows系统API
typedef unsigned char u_char;
typedef unsigned short u_short;
typedef unsigned int u_int;
typedef unsigned long u_long;
#endif

#include <pcap.h>

using namespace std;

// FTP登录结构体
struct FTPLogin {
    string timestamp;
    string src_mac;
    string src_ip;
    string dst_mac;
    string dst_ip;
    string username;
    string password;
    string success;
};

// TCP流结构体
struct TcpStream {
    string buffer;  //收集payload字符串
    u_int32_t next_seq;  //下一个应收的TCP seq
    bool is_active;  //标记是否为活跃连接
};

unordered_map<string, TcpStream> tcp_streams;
vector<FTPLogin> ftp_logins;

// MACַ格式转换
string format_mac(const u_char* mac) {
    char buf[18];
    sprintf_s(buf, "%02X-%02X-%02X-%02X-%02X-%02X",
        mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
    return string(buf);
}

// IPַ地址格式转换
string format_ip(u_long addr) {
    struct in_addr ip_addr;
    ip_addr.s_addr = addr;
    char buf[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &ip_addr, buf, INET_ADDRSTRLEN);
    return string(buf);
}

//时间格式转换
string format_timestamp(timeval ts) {
    time_t time = ts.tv_sec;
    struct tm tm_info;
    localtime_s(&tm_info, &time);
    char buf[20];
    strftime(buf, 20, "%Y-%m-%d %H:%M:%S", &tm_info);
    return string(buf);
}

//协议头格式定义
#pragma pack(push, 1)
struct ether_header {
    u_char ether_dhost[6];
    u_char ether_shost[6];
    u_short ether_type;
};

struct ip_header {
    u_char  ihl : 4;
    u_char  version : 4;
    u_char  tos;
    u_short tot_len;
    u_short id;
    u_short frag_off;
    u_char  ttl;
    u_char  protocol;
    u_short check;
    u_long  saddr;
    u_long  daddr;
};

struct tcp_header {
    u_short sport;
    u_short dport;
    u_int   seq;
    u_int   ack;
    u_char  off_res;
    u_char  flags;
    u_short win;
    u_short check;
    u_short urg_ptr;
};
#pragma pack(pop)

//分析一帧捕获到的网络数据包，提取 FTP 登录相关信息
void parse_packet(u_char* user_data, const pcap_pkthdr* pkthdr, const u_char* packet, int link_type) {
    const u_char* network_layer = packet;
    size_t link_header_size = 0;
    //判断链路层头大小 若遇到不支持的类型就跳过处理
    switch (link_type) {
    case 0:
        link_header_size = 4;
        break;
    case 1:
        link_header_size = sizeof(ether_header);
        break;
    default:
        cerr << "Unsupported link type: " << link_type << endl;
        return;
    }
    //获取 IP 头并过滤非 IPv4、非 TCP 包
    network_layer += link_header_size;
    ip_header* iph = (ip_header*)network_layer;
    if (iph->version != 4) return;
    if (iph->protocol != IPPROTO_TCP) return;
    //验证 IP 头长度，定位 TCP 头
    u_int ip_len = iph->ihl * 4;
    if (ip_len < 20) {
        cerr << "Invalid IP header length: " << iph->ihl << endl;
        return;
    }
    tcp_header* tcph = (tcp_header*)(network_layer + ip_len);
    //过滤非 FTP 会话包（不是端口21）
    u_short src_port = ntohs(tcph->sport);
    u_short dst_port = ntohs(tcph->dport);
    if (dst_port != 21 && src_port != 21) return;
    //构建 TCP 流标识符 & 初始化流状态
    string stream_id = format_ip(iph->saddr) + ":" + to_string(src_port) + "-" +
        format_ip(iph->daddr) + ":" + to_string(dst_port);

    if (!tcp_streams.count(stream_id)) {
        tcp_streams[stream_id] = TcpStream{ "", ntohl(tcph->seq), true };
    }
    //提取 TCP Payload（FTP 命令字符串）
    u_int tcp_off = (tcph->off_res >> 4) * 4;
    if (tcp_off < 20) return;
    const u_char* payload = network_layer + ip_len + tcp_off;
    int payload_len = pkthdr->caplen - (link_header_size + ip_len + tcp_off);
    if (payload_len <= 0) return;
    tcp_streams[stream_id].buffer.append((char*)payload, payload_len);
    //处理 FTP 命令行（逐行解析）
    string& buffer = tcp_streams[stream_id].buffer;
    size_t pos;
    while ((pos = buffer.find("\r\n")) != string::npos) {
        string line = buffer.substr(0, pos);
        buffer.erase(0, pos + 2);
        cerr << "FTP Line: " << line << endl;
        //提取 USER 命令并构造登录记录
        if (line.find("USER ") == 0) {
            FTPLogin login;
            login.timestamp = format_timestamp(pkthdr->ts);

            if (link_type == 1) {
                ether_header* eth = (ether_header*)packet;
                login.src_mac = format_mac(eth->ether_shost);
                login.dst_mac = format_mac(eth->ether_dhost);
            }
            else {
                login.src_mac = "00-00-00-00-00-00";
                login.dst_mac = "00-00-00-00-00-00";
            }

            login.src_ip = format_ip(iph->saddr);
            login.dst_ip = format_ip(iph->daddr);

            size_t start = line.find("USER ") + 5;
            login.username = line.substr(start);
            ftp_logins.push_back(login);
        }
        //提取 PASS 密码
        else if (line.find("PASS ") == 0 && !ftp_logins.empty()) {
            size_t start = line.find("PASS ") + 5;
            ftp_logins.back().password = line.substr(start);
        }
        //判断是否登录成功/失败
        else if (isdigit(line[0]) && line.size() >= 3) {
            string code = line.substr(0, 3);
            if (!ftp_logins.empty()) {
                if (code == "230") {
                    ftp_logins.back().success = "SUCCEED";
                    cerr << "Login succeeded for user: " << ftp_logins.back().username << endl;
                }
                else if (code == "530") {
                    ftp_logins.back().success = "FAILED";
                }
                else if (code == "331") {
                    cerr << "Server requested password" << endl;
                }
            }
        }
    }
    //清理不活跃连接
    static int packet_count = 0;
    if (++packet_count % 100 == 0) {
        for (auto it = tcp_streams.begin(); it != tcp_streams.end(); ) {
            if (!it->second.is_active) {
                cerr << "Cleaning inactive stream: " << it->first << endl;
                it = tcp_streams.erase(it);
            }
            else {
                ++it;
            }
        }
    }
}

int main(int argc, char* argv[]) {
    if (argc != 3) {
        cerr << "Usage: " << argv[0] << " <input.pcapng> <output.csv>" << endl;
        return 1;
    }

    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_t* pcap = pcap_open_offline(argv[1], errbuf);
    if (!pcap) {
        cerr << "Error opening pcap file: " << errbuf << endl;
        return 1;
    }

    int link_type = pcap_datalink(pcap);
    cerr << "Link layer type: " << link_type << endl;

    if (pcap_set_timeout(pcap, 1000) < 0) {
        cerr << "Warning: Set timeout failed" << endl;
    }

    auto packet_handler = [&](u_char* user, const pcap_pkthdr* h, const u_char* p) {
        parse_packet(user, h, p, link_type);
        };

    if (pcap_loop(pcap, 0, [](u_char* user, const pcap_pkthdr* h, const u_char* p) {
        auto handler = reinterpret_cast<decltype(packet_handler)*>(user);
        (*handler)(user + sizeof(handler), h, p);
        }, reinterpret_cast<u_char*>(&packet_handler)) < 0) {
        cerr << "Error processing pcap file" << endl;
        pcap_close(pcap);
        return 1;
    }

    cerr << "Found " << ftp_logins.size() << " FTP logins." << endl;
    pcap_close(pcap);

    ofstream csv(argv[2], ios::out | ios::binary);
    if (!csv) {
        cerr << "Error opening output file" << endl;
        return 1;
    }

    csv << "ʱ,ԴMAC,ԴIP,ĿMAC,ĿIP,¼,,ɹ\n";
    for (const auto& login : ftp_logins) {
        csv << login.timestamp << ","
            << login.src_mac << ","
            << login.src_ip << ","
            << login.dst_mac << ","
            << login.dst_ip << ","
            << login.username << ","
            << login.password << ","
            << login.success << "\n";
    }
    csv.flush();
    csv.close();

    if (!csv.good()) {
        cerr << "Error writing output file" << endl;
        return 1;
    }

    cout << "Processed " << ftp_logins.size() << " FTP logins. Output saved to " << argv[2] << endl;
    return 0;
}